import numpy as np
import re
from collections import defaultdict, Counter
from math import log

class NaiveBayesSpamDetector:
    def __init__(self):
        self.word_counts = defaultdict(lambda: defaultdict(int))
        self.class_counts = defaultdict(int)
        self.vocabulary = set()
        self.classes = []
    
    def preprocess_text(self, text):
        """文本预处理：转小写，提取单词"""
        text = text.lower()
        words = re.findall(r'\b\w+\b', text)
        return words
    
    def fit(self, emails, labels):
        """训练模型"""
        self.classes = list(set(labels))
        
        for email, label in zip(emails, labels):
            words = self.preprocess_text(email)
            self.class_counts[label] += 1
            
            for word in words:
                self.word_counts[label][word] += 1
                self.vocabulary.add(word)
    
    def predict_proba(self, email):
        """预测概率"""
        words = self.preprocess_text(email)
        class_probabilities = {}
        
        total_emails = sum(self.class_counts.values())
        
        for class_name in self.classes:
            # 计算先验概率
            class_prob = log(self.class_counts[class_name] / total_emails)
            
            # 计算似然概率（拉普拉斯平滑）
            total_words_in_class = sum(self.word_counts[class_name].values())
            vocab_size = len(self.vocabulary)
            
            for word in words:
                word_count = self.word_counts[class_name][word]
                # 拉普拉斯平滑：避免零概率
                word_prob = log((word_count + 1) / (total_words_in_class + vocab_size))
                class_prob += word_prob
            
            class_probabilities[class_name] = class_prob
        
        return class_probabilities
    
    def predict(self, email):
        """预测类别"""
        probabilities = self.predict_proba(email)
        return max(probabilities.items(), key=lambda x: x[1])[0]

# 测试垃圾邮件检测器
if __name__ == "__main__":
    # 训练数据
    emails = [
        "Free money! Click here now! Limited time offer!",
        "Hi John, let's meet for coffee tomorrow at 3pm",
        "URGENT: Your account will be closed! Send money now!",
        "Thanks for the meeting notes, I'll review them tonight",
        "Congratulations! You've won a million dollars! Click here!",
        "Can you send me the project report by Friday?",
        "WINNER! Claim your prize now! No payment required!",
        "Let's discuss the budget in tomorrow's meeting"
    ]
    
    labels = ['spam', 'ham', 'spam', 'ham', 'spam', 'ham', 'spam', 'ham']
    
    # 训练模型
    detector = NaiveBayesSpamDetector()
    detector.fit(emails, labels)
    
    # 测试
    test_emails = [
        "Free offer! Click now!",
        "Meeting at 2pm tomorrow",
        "You won! Send money!"
    ]
    
    for email in test_emails:
        prediction = detector.predict(email)
        probabilities = detector.predict_proba(email)
        print(f"Email: {email}")
        print(f"Prediction: {prediction}")
        print(f"Probabilities: {probabilities}")
        print("-" * 50) 

        